package com.fzs.sshd.cli.server.shell;

import com.fzs.sshd.cli.server.session.SessionAttributeKey;
import org.apache.sshd.server.Environment;
import org.apache.sshd.server.ExitCallback;
import org.apache.sshd.server.channel.ChannelSession;
import org.apache.sshd.server.command.Command;
import org.apache.sshd.server.shell.ShellFactory;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.atomic.AtomicBoolean;

public class ADBWebShellProcessShellFactory implements ShellFactory {

    private static final String ADB_TARGET = "183.240.23.14:37223";
    private static final String ANSI_ESCAPE_PATTERN = "\\u001B\\[([0-9]{1,3}(;[0-9]{1,3})*)?[mK]";

    public ADBWebShellProcessShellFactory() {}

    public static String removeAnsiCodes(String input) {
        if (input == null) return "";
        return input.replaceAll(ANSI_ESCAPE_PATTERN, "");
    }

    @Override
    public Command createShell(ChannelSession session) {
        return new ADBCommand(session, ADB_TARGET);
    }

    private static class ADBCommand implements Command {

        private final ChannelSession session;
        private final String adbTarget;
        private final AtomicBoolean closed = new AtomicBoolean(false);

        private InputStream in;
        private OutputStream out;
        private OutputStream err;
        private Process shellProcess;

        public ADBCommand(ChannelSession session, String adbTarget) {
            this.session = session;
            this.adbTarget = adbTarget;
        }

        @Override
        public void start(ChannelSession channelSession, Environment environment) {

            try {
                // 1️⃣ 欢迎信息
                writeLine(out, "");
                writeLine(out, "📱 Welcome to ADB Remote Console!");
                writeLine(out, "🔌 Automatically connecting to: " + adbTarget);
                writeLine(out, "💡 Type 'exit' or 'quit' to disconnect.");
                writeLine(out, "🚫 No shell access allowed. Only ADB commands are permitted.");
                writeLine(out, "⏰ Session force close in " +
                        channelSession.getSession().getAttribute(SessionAttributeKey.EXPIRE_TIME));
                writeLine(out, "");

                // 2️⃣ adb connect
                boolean alreadyConnected = false;
                ProcessBuilder pbConnect = new ProcessBuilder("adb", "connect", adbTarget);
                pbConnect.redirectErrorStream(true);
                Process connectProcess = pbConnect.start();
                try (BufferedReader reader = new BufferedReader(
                        new InputStreamReader(connectProcess.getInputStream(), StandardCharsets.UTF_8))) {
                    String line;
                    while ((line = reader.readLine()) != null) {
                        String clean = removeAnsiCodes(line);
                        write(out, clean + "\r\n");
                        if (line.contains("already connected")) {
                            alreadyConnected = true;
                        }
                    }
                }
                connectProcess.waitFor();

                if (!alreadyConnected) {
                    writeLine(out, "✅ Connected to " + adbTarget);
                }

                // 3️⃣ 启动 adb shell
                ProcessBuilder pbShell = new ProcessBuilder("adb", "shell");
                pbShell.redirectErrorStream(true);
                shellProcess = pbShell.start();
                final Process process = shellProcess;

                // 4️⃣ 输出转发线程
                Thread outputForwarder = new Thread(() -> forwardProcessOutput(process));
                outputForwarder.setDaemon(true);
                outputForwarder.start();

                // 5️⃣ 输入转发线程
                Thread inputForwarder = new Thread(() -> forwardClientInput(process));
                inputForwarder.setDaemon(true);
                inputForwarder.start();

            } catch (IOException | InterruptedException e) {
                closeSession();
            }
        }

        private void forwardProcessOutput(Process process) {
            try (InputStream is = process.getInputStream()) {
                byte[] buf = new byte[1024];
                int n;
                while (!closed.get() && (n = is.read(buf)) > 0) {
                    String clean = removeAnsiCodes(new String(buf, 0, n, StandardCharsets.UTF_8));
                    out.write(clean.getBytes(StandardCharsets.UTF_8));
                    out.flush();
                }
            } catch (IOException ignored) {}
        }

        private void forwardClientInput(Process process) {
            try (OutputStream processIn = process.getOutputStream()) {
                byte[] buf = new byte[1024];
                int len;
                while (!closed.get() && (len = in.read(buf)) > 0) {
                    String input = new String(buf, 0, len, StandardCharsets.UTF_8);
                    if ("exit".equalsIgnoreCase(input.trim()) || "quit".equalsIgnoreCase(input.trim())) {
                        writeLine(out, "👋 Disconnecting from ADB console...");
                        destroy(session);
                        return;
                    }
                    processIn.write(buf, 0, len);
                    processIn.flush();
                }
            } catch (IOException ignored) {}
        }

        @Override
        public void destroy(ChannelSession channelSession) {
            if (closed.getAndSet(true)) return;

            try {
                if (shellProcess != null && shellProcess.isAlive()) {
                    shellProcess.destroyForcibly();
                }

                // adb disconnect
                ProcessBuilder pbDisconnect = new ProcessBuilder("adb", "disconnect", adbTarget);
                pbDisconnect.redirectErrorStream(true);
                pbDisconnect.start();

            } catch (IOException ignored) {}
            closeSession();
        }

        private void closeSession() {
            if (session != null && session.isOpen()) {
                session.close(false);
            }
        }

        private static void writeLine(OutputStream out, String text) throws IOException {
            out.write((text + "\r\n").getBytes(StandardCharsets.UTF_8));
            out.flush();
        }

        private static void write(OutputStream out, String text) throws IOException {
            out.write(text.getBytes(StandardCharsets.UTF_8));
            out.flush();
        }

        @Override
        public void setInputStream(InputStream in) { this.in = in; }

        @Override
        public void setOutputStream(OutputStream out) { this.out = out; }

        @Override
        public void setErrorStream(OutputStream err) { this.err = err; }

        @Override
        public void setExitCallback(ExitCallback callback) {}
    }
}
